import math
import numpy as np
import spconv.pytorch as spconv
import torch
from torch import nn
from functools import partial
from einops import rearrange

from DiffOGMP.utils.common_utils import instantiate_from_config

class BindedBN1d(nn.Module):
    def __init__(self, planes, eps=1e-3, momentum=0.01):
        super().__init__()
        self.bn = nn.BatchNorm1d(planes, eps=eps, momentum=momentum)

    def forward(self, x):
        return self.bn(x)


def post_act_block(in_channels, out_channels, kernel_size, indice_key=None, stride=1, padding=0,
                   conv_type='subm', norm_fn=None):

    if conv_type == 'subm':
        conv = spconv.SubMConv3d(in_channels, out_channels, kernel_size, bias=False, indice_key=indice_key)
    elif conv_type == 'spconv':
        conv = spconv.SparseConv3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding,
                                   bias=False, indice_key=indice_key)
    elif conv_type == 'inverseconv':
        conv = spconv.SparseInverseConv3d(in_channels, out_channels, kernel_size, indice_key=indice_key, bias=False)
    else:
        raise NotImplementedError

    m = spconv.SparseSequential(
        conv,
        norm_fn(out_channels),
        nn.ReLU(),
    )

    return m

def replace_feature(out, new_features):
    if "replace_feature" in out.__dir__():
        # spconv 2.x behaviour
        return out.replace_feature(new_features)
    else:
        out.features = new_features
        return out

class SparseBasicBlock(spconv.SparseModule):
    def __init__(self,
                 in_channels,
                 channels,
                 stride=1,
                 norm_fn=None,
                 downsample=None,
                 indice_key=None):
        super().__init__()
        assert norm_fn is not None
        bias = norm_fn is not None
        self.conv1 = spconv.SubMConv3d(
            in_channels, channels, kernel_size=3, stride=stride, padding=1, bias=bias, indice_key=indice_key
        )
        self.bn1 = norm_fn(channels)
        self.relu = nn.ReLU()
        self.conv2 = spconv.SubMConv3d(
            channels, channels, kernel_size=3, stride=stride, padding=1, bias=bias, indice_key=indice_key
        )
        self.bn2 = norm_fn(channels)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        identity = x
        out = self.conv1(x)
        out = replace_feature(out, self.bn1(out.features))
        out = replace_feature(out, self.relu(out.features))

        out = self.conv2(out)
        out = replace_feature(out, self.bn2(out.features))

        if self.downsample is not None:
            identity = self.downsample(x)
        out = replace_feature(out, out.features + identity.features)
        out = replace_feature(out, self.relu(out.features))

        return out


class AttnBlock(nn.Module):
    def __init__(self, in_channels, num_groups=8):
        super().__init__()
        self.in_channels = in_channels

        self.norm = nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
        self.q = nn.Conv3d(in_channels,
                             in_channels,
                             kernel_size=1,
                             stride=1,
                             padding=0)
        self.k = nn.Conv3d(in_channels,
                             in_channels,
                             kernel_size=1,
                             stride=1,
                             padding=0)
        self.v = nn.Conv3d(in_channels,
                             in_channels,
                             kernel_size=1,
                             stride=1,
                             padding=0)
        self.proj_out = nn.Conv3d(in_channels,
                                    in_channels,
                                    kernel_size=1,
                                    stride=1,
                                    padding=0)


    def forward(self, x):
        if isinstance(x, spconv.SparseConvTensor):
            h_ = x.dense()
        else:
            h_ = x
        b, c, t, h, w = h_.size()
        h_ = self.norm(h_)
        q = self.q(h_)
        k = self.k(h_)
        v = self.v(h_)

        # compute attention
        q = rearrange(q, 'b c t h w -> b (t h w) c')
        k = rearrange(k, 'b c t h w -> b c (t h w)')
        v = rearrange(v, 'b c t h w -> b c (t h w)')
        w_ = torch.einsum('b i c, b c j -> b i j', q, k)
        w_ = w_ * (int(c)**(-0.5))
        w_ = torch.nn.functional.softmax(w_, dim=2)
        h_ = torch.einsum('b i j, b c j -> b i c', w_, v)
        h_ = rearrange(h_, 'b (t h w) c -> b c t h w', t=t, h=h, w=w)

        h_ = self.proj_out(h_)

        return x+h_


class Encoder3d(spconv.SparseModule):
    def __init__(self, z_channels, double_z, channels=(16, 32, 64, 128), image_channel=1):
        super().__init__()
        norm_fn = partial(nn.BatchNorm1d, eps=1e-3, momentum=0.01)
        block = post_act_block
        self.conv_first = spconv.SparseSequential(
            spconv.SubMConv3d(image_channel, channels[0], 3, 1, indice_key='subm1'),
            norm_fn(channels[0]),
            nn.ReLU())
        self.conv1 = spconv.SparseSequential(
            SparseBasicBlock(channels[0], channels[0], norm_fn=norm_fn, indice_key='res1'),
            SparseBasicBlock(channels[0], channels[0], norm_fn=norm_fn, indice_key='res1')
        )
        self.conv2 = spconv.SparseSequential(
            block(channels[0], channels[1], 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv1', conv_type='spconv'),
            SparseBasicBlock(channels[1], channels[1], norm_fn=norm_fn, indice_key='res2'),
            SparseBasicBlock(channels[1], channels[1], norm_fn=norm_fn, indice_key='res2')
        )
        self.conv3 = spconv.SparseSequential(
            block(channels[1], channels[2], 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv2', conv_type='spconv'),
            SparseBasicBlock(channels[2], channels[2], norm_fn=norm_fn, indice_key='res3'),
            SparseBasicBlock(channels[2], channels[2], norm_fn=norm_fn, indice_key='res3')
        )
        self.conv4 = spconv.SparseSequential(
            block(channels[2], channels[3], 3, norm_fn=norm_fn, stride=(1,2,2), padding=1, indice_key='spconv3', conv_type='spconv'),
            SparseBasicBlock(channels[3], channels[3], norm_fn=norm_fn, indice_key='res4'),
            SparseBasicBlock(channels[3], channels[3], norm_fn=norm_fn, indice_key='res4'),
            spconv.ToDense()
        )

        self.mid = AttnBlock(channels[3])
        self.conv_final = nn.Sequential(
            nn.GroupNorm(8, channels[3], eps=1e-6, affine=True),
            nn.SiLU(),
            nn.Conv3d(channels[3], 2 * z_channels if double_z else z_channels, 3, 1, 1)
        )

    def forward(self, x):
        x_sp = spconv.SparseConvTensor.from_dense(x.permute(0, 1, 3, 4, 2))
        out = self.conv_first(x_sp)
        out = self.conv1(out)
        out = self.conv2(out)
        out = self.conv3(out)
        out = self.conv4(out)
        out = self.mid(out)
        out = self.conv_final(out)
        return out

class Decoder3d(spconv.SparseModule):
    def __init__(self, z_channels, double_z, channels=(128, 64, 32, 16), image_channel=1):
        super().__init__()
        norm_fn = partial(nn.BatchNorm1d, eps=1e-3, momentum=0.01)
        block = post_act_block
        self.conv_first = spconv.SparseSequential(
            spconv.SubMConv3d(z_channels, channels[0], 3, 1, 1, indice_key='conv_first_dec'),
            norm_fn(channels[0]),
            nn.ReLU())
        self.conv1 = spconv.SparseSequential(
            SparseBasicBlock(channels[0], channels[0], norm_fn=norm_fn, indice_key='res1'),
            SparseBasicBlock(channels[0], channels[0], norm_fn=norm_fn, indice_key='res1')
        )
        self.conv2 = spconv.SparseSequential(
            block(channels[0], channels[1], 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv3',
                  conv_type='inverseconv'),
            SparseBasicBlock(channels[1], channels[1], norm_fn=norm_fn, indice_key='res2'),
            SparseBasicBlock(channels[1], channels[1], norm_fn=norm_fn, indice_key='res2')
        )
        self.conv3 = spconv.SparseSequential(
            block(channels[1], channels[2], 3, norm_fn=norm_fn, stride=1, padding=1, indice_key='spconv2',
                  conv_type='inverseconv'),
            SparseBasicBlock(channels[2], channels[2], norm_fn=norm_fn, indice_key='res3'),
            SparseBasicBlock(channels[2], channels[2], norm_fn=norm_fn, indice_key='res3')
        )
        self.conv4 = spconv.SparseSequential(
            block(channels[2], channels[3], 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv1',
                  conv_type='inverseconv'),
            SparseBasicBlock(channels[3], channels[3], norm_fn=norm_fn, indice_key='res4'),
            SparseBasicBlock(channels[3], channels[3], norm_fn=norm_fn, indice_key='res4'),
            spconv.ToDense()
        )
        self.conv_final = nn.Conv3d(channels[3], image_channel, 3, 1,1)

    def forward(self, x):
        # x -> [N, C, T, H, W]
        if isinstance(x, spconv.SparseConvTensor):
            x_sp = x
        else:
            x_sp = spconv.SparseConvTensor.from_dense(x.permute(0, 2, 3, 4, 1))
        out = self.conv_first(x_sp)
        out = self.conv1(out)
        out = self.conv2(out)
        out = self.conv3(out)
        out = self.conv4(out)
        out = self.conv_final(out)
        return out